
#include <stdio.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <stdlib.h>
#include <sys/time.h>
#include <signal.h>
#include <sys/wait.h>

#include "../../locking.h"
#include "../../sr_module.h"
#include "../../dprint.h"
#include "../../error.h"
#include "../../ut.h"
#include "../../mem/mem.h"
#include "../../mem/shm_mem.h"
#include "../../str.h"
#include "../../parser/parse_uri.h"
#include "../../parser/parse_from.h"
#include "../../parser/parse_to.h"
#include "../../modules/tm/tm_load.h"
#include "../dialog/dlg_load.h"
#include "../dialog/dlg_hash.h"
#include "../../modules/tm/tm_load.h"
#include "../../rpc_lookup.h"

#include "hash.h"


/* Required in every Kamailio Module. */
MODULE_VERSION


/*! This is the first function to be called by Kamailio, to initialize the module.
 * This call must always return a value as soon as possible.  If it were not to
 * return, then Kamailio would not be able to initialize any of the other
 * modules. */
static int  mod_init(void);

/*! This function is called when Kamailio has finished creating all instances of
 * itself.  It is at this point that we want to create our AgentX sub-agent
 * process, and register a handler for any state changes of our child. */
static int  mod_child_init(int rank);


/*! This function is called when Kamailio is shutting down.  When this happens, we
 * log a useful message and kill the AgentX Sub-Agent child process */
static void mod_destroy(void);





/**********************
 *  local functions
 **********************/
static int init_dialog_callbacks(void);

static void on_dialog_create(struct dlg_cell *dlg, int type, struct dlg_cb_params *_params);

static void on_dialog_update(struct dlg_cell *dlg, int type, struct dlg_cb_params *_params);

static void handle_callback(struct dlg_cell *dlg, int type);

struct dlg_binds dialog_api;

struct tm_binds tm_api;

/*
 * The number of the maximum allowed cross-proxy dialogs.
 * If set to zero, no limitation will be applied.
 * This value is overridden by the exported parameter which is set in configure file.
 */
int maxDialogNumber = 4;

/*
 * This is where we keep the records of every proxy's cross-proxy dialog amount.
 */
HashTable *maxDialogHashtable;

/***********************
 *  exported functions
 ***********************/
static int w_check_max_dialog(struct sip_msg*, char*, char*);


/*
 * The exported commands
 */
static cmd_export_t cmds[]= {
    {
        "check_max_dialog", (cmd_function)w_check_max_dialog, 0, 0, 0, REQUEST_ROUTE
    },
    {0,0,0,0,0,0}
};

/*!
 * This structure defines the SNMPStats parameters that can be configured
 * through the kamailio.cfg configuration file.
 */
static param_export_t mod_params[] = {
    { "maxDialogs", INT_PARAM, &maxDialogNumber},
    { 0,0,0 }
};

/*
 *  The RPC declaration
 */
static void rpc_print_table(rpc_t *rpc, void *c);

static const char *rpc_print_table_doc[2] = {
	"Print all proxy's active dialog number", 0
};

static rpc_export_t rpc_methods[] = {
	{"quantum.list", rpc_print_table, rpc_print_table_doc, 0},
	{0, 0, 0, 0}
};


struct module_exports exports = {
    "quantum_control",       /* module's name */
    DEFAULT_DLFLAGS,         /* dlopen flags */
    cmds,                    /* exported functions */
    mod_params,              /* param exports */
    0,                       /* exported statistics */
    0,                       /* MI Functions */
    0,                       /* pseudo-variables */
    0,                       /* extra processes */
    mod_init,                /* module initialization function */
    0,                       /* reply processing function */
    mod_destroy,             /* Destroy function */
    mod_child_init           /* per-child init function */
};



/*! This is the first function to be called by Kamailio, to initialize the module.
 * This call must always return a value as soon as possible.  If it were not to
 * return, then Kamailio would not be able to initialize any of the other
 * modules. */
static int mod_init(void) {

    INFO("Initializing Quantum Control Module...");
    INFO("Maximum allowed dialogs set to %d", maxDialogNumber);

    init_dialog_callbacks();

    // init hashtable
    maxDialogHashtable = hash_create(1024);
    if (!maxDialogHashtable) {
        LM_ERR("Cannot create hashtable\n");
        return -1;
    }

    // register RPC
    if (0 != rpc_register_array(rpc_methods)) {
		LM_ERR("failed to register RPC commands\n");
		return -1;
	}

    return 0;
}


/*! This function is called when Kamailio has finished creating all instances of
 * itself.
 */
static int mod_child_init(int rank) {
    return 0;
}

/*! This function is called when Kamailio is shutting down. When this happens, we
 * log a useful message and kill the AgentX Sub-Agent child process */
static void mod_destroy(void) {
    hash_free(maxDialogHashtable);
}




int init_dialog_callbacks() {

    /* bind to the dialog API */
    if (0 != load_dlg_api(&dialog_api)) {
        LM_ERR("failed to load dialog API - is dialog module loaded?\n");
        return -1;
    }

    /* register dialog creation callback */
    if (0 != dialog_api.register_dlgcb(NULL, DLGCB_CREATED, on_dialog_create, NULL, NULL)) {
        LM_ERR("Cannot register callback to dialog creation\n");
        return -1;
    }

    return 0;
}

void on_dialog_create(struct dlg_cell *dlg, int type, struct dlg_cb_params *_params) {

    if (0 != load_dlg_api(&dialog_api)) { //Must reload dialog api here.
        LM_ERR("failed to load dialog API - is dialog module loaded?\n");
        return;
    }

    // register all per-dialog callbacks except for DLGCB_DESTROY
    if (0 != dialog_api.register_dlgcb(dlg,
                                       DLGCB_FAILED | DLGCB_TERMINATED | DLGCB_EXPIRED,
                                       on_dialog_update, NULL, NULL)) {
        LM_ERR("Cannot register callbacks for per-dialog callbacks\n");
        return;
    }

    /*
     struct sip_msg *msg;
     msg = dlg_get_valid_msg(_params);

     load_tm_api(&tm_api);
     tm_api.t_newtran(msg);
     tm_api.t_reply(msg, 507, "TOO MANY CALLS 2");
     */
}

void on_dialog_update(struct dlg_cell *dlg, int type, struct dlg_cb_params *_params) {
    handle_callback(dlg, type);
}

void handle_callback(struct dlg_cell *dlg, int type) {
    /*
    LM_INFO("QC call back handler. Type=>");
    switch (type) {
    case DLGCB_CREATED:
        INFO("DLGCB_CREATED");
        break;
    case DLGCB_EARLY:
        INFO("DLGCB_EARLY");
        break;
    case DLGCB_FAILED:
        INFO("DLGCB_FAILED");
        break;
    case DLGCB_CONFIRMED_NA:
        INFO("DLGCB_CONFIRMED_NA");
        break;
    case DLGCB_CONFIRMED:
        INFO("DLGCB_CONFIRMED");
        break;
    case DLGCB_REQ_WITHIN:
        INFO("DLGCB_REQ_WITHIN");
        break;
    case DLGCB_TERMINATED:
        INFO("DLGCB_TERMINATED");
        break;
    case DLGCB_EXPIRED:
        INFO("DLGCB_EXPIRED");
        break;
    case DLGCB_RESPONSE_FWDED:
        INFO("DLGCB_RESPONSE_FWDED");
        break;
    case DLGCB_RESPONSE_WITHIN:
        INFO("DLGCB_RESPONSE_WITHIN");
        break;
    case DLGCB_MI_CONTEXT:
        INFO("DLGCB_MI_CONTEXT");
        break;
    case DLGCB_DESTROY:
        INFO("DLGCB_DESTROY");
        break;
    default:
        LM_ERR("dialog callback type 'unknown' received, from=%.*s\n", dlg->from_uri.len, dlg->from_uri.s);
    }
    */

    // Parse quantumId
    // call id
    char* callId = pkg_malloc(dlg->callid.len * sizeof(char) + 1);
    strncpy(callId, dlg->callid.s, dlg->callid.len);
    callId[dlg->callid.len] = '\0';

    // Parse calleeId from 'to_uri'
    char temp[100];
    strncpy(temp, dlg->to_uri.s, dlg->to_uri.len);
    temp[dlg->to_uri.len] = '\0';

    strtok(temp, ":");

    char *tmpString = strtok(NULL, "@");
    char *calleeId = malloc(strlen(tmpString) * sizeof(char) + 1);
    strcpy(calleeId, tmpString);
    calleeId[strlen(tmpString)] = '\0';

    // Parse other data from 'from_uri'
    memset(temp, 0, sizeof(temp));
    strncpy(temp, (const char*)(dlg->from_uri.s), dlg->from_uri.len);

    // The string should be like this:
    // sip:callerId#callerProxyId#calleeProxyId#param@serverAddress
    strtok(temp, ":");
    /*char *callerId = */strtok(NULL, "#");
    char *callerProxyId = strtok(NULL, "#");
    char *calleeProxyId = strtok(NULL, "#");

    // Skip the calls from the same quantum id.
    if (callerProxyId && calleeProxyId && strcmp(callerProxyId, calleeProxyId) == 0)
        return;

    // Skip if qid doesn't exist
    if (!callerProxyId || !calleeProxyId)
        return;

    if (type == DLGCB_FAILED || type == DLGCB_TERMINATED || type == DLGCB_EXPIRED) {
        int *ptr;
        ptr = (int *)hash_get(maxDialogHashtable, callerProxyId);
        if (ptr && *ptr > 0)
            (*ptr)--;

        ptr = (int *)hash_get(maxDialogHashtable, calleeProxyId);
        if (ptr && *ptr > 0)
            (*ptr)--;

        //hash_print(maxDialogHashtable);
    }
}

/*
 * This function is called in the RouteLogic right before calling async route to RS.
 * It checks the number of cross-proxy dialogs related to this INVITE. either caller
 * or callee's proxy reaching the threshold will cause a 506 reply. From the aspect
 * of dialog module, even the INVITE is rejected, the dialog is still created and
 * then failed.
 */
static int w_check_max_dialog(struct sip_msg *sip_msg, char *foo, char *bar) {
    INFO("Checking quantum proxy's dialog threshold......");

    // Always validate the check if maxDialogNumber is set to 0.
    if (maxDialogNumber == 0)
        return 1;

    // Validate for RE-INVITE requests
    str tag;
    if (!sip_msg->to) {
        // to header not defined, parse to header
        LM_DBG("to header not defined, parse to header\n");
        if (parse_headers(sip_msg, HDR_TO_F,0) < 0) {
            LM_ERR("parsing of to-header failed\n");
            tag.s = 0;
            tag.len = 0;
        } else if (!sip_msg->to) {
            LM_DBG("no to although to-header is parsed: bad reply or missing TO hdr :-/\n");
            tag.s = 0;
            tag.len = 0;
        } else
            tag = get_to(sip_msg)->tag_value;
    } else {
        tag = get_to(sip_msg)->tag_value;
        if (tag.s==0 || tag.len==0) {
            LM_DBG("missing TAG param in TO hdr :-/\n");
            tag.s = 0;
            tag.len = 0;
        }
    }
    if (tag.s) {
        LM_INFO("RE-INVITE detected. to_tag=[%.*s]", tag.len, tag.s);
        return 1;
    }

    // Parse the caller & callee's quantum id
    char uri[get_from(sip_msg)->uri.len + 1];
    strncpy(uri, get_from(sip_msg)->uri.s, get_from(sip_msg)->uri.len);
    uri[get_from(sip_msg)->uri.len] = '\0';
    //INFO("Parsing [%s]", uri);

    strtok(uri, ":");
    strtok(NULL, "#");
    char *callerQid = strtok(NULL, "#");
    char *calleeQid = strtok(NULL, "#");

    // Validate the check in case qid is missing.
    if (!callerQid || !calleeQid)
        return 1;

    // Validate the check if two qid are equal.
    if (strcmp(callerQid, calleeQid) == 0)
        return 1;

    INFO("Qid: %s -> %s", callerQid, calleeQid);

    // Lock the hashtable
    lock_get(maxDialogHashtable->lock);

    int result;
    int *callerNum;
    int *calleeNum;

    callerNum = (int *)hash_get(maxDialogHashtable, callerQid);
    if (!callerNum) {
        callerNum = shm_malloc(sizeof(int));
        *callerNum = 0;
        hash_put(maxDialogHashtable, callerQid, callerNum);
    }

    calleeNum = (int *)hash_get(maxDialogHashtable, calleeQid);
    if (!calleeNum) {
        calleeNum = shm_malloc(sizeof(int));
        *calleeNum = 0;
        hash_put(maxDialogHashtable, calleeQid, calleeNum);
    }

    result = ( *callerNum >= maxDialogNumber || *calleeNum >= maxDialogNumber) ? -1 : 1;

    if (result == -1){
        INFO("check_max_dialog failed. Max allowed is %d", maxDialogNumber);
        hash_print(maxDialogHashtable);
    }

    // Always ++ here regardless the result because no matter what the result is, a dialog
    // callback indicating end of the dialog will be called and these two numbers will
    // decrease by 1 by then.
    (*callerNum)++;
    (*calleeNum)++;

    //hash_print(maxDialogHashtable);

    lock_release(maxDialogHashtable->lock);

    return result;
}

/*
 * The RPC command function which prints all key-value pairs in 'maxDialogHashtable'
 */
static void rpc_print_table(rpc_t *rpc, void *c) {
    unsigned int size = hash_getSize(maxDialogHashtable);
    int i;
    for (i = 0; i < size; i++) {
        HashNode *entry = maxDialogHashtable->ht[i];
        while (entry) {
            rpc->printf(c, "[%s] -> %d", entry->key, *((int *)(entry->value)));
            entry = entry->next;
        }
    }
}
